package cn.edu.flink.scala.tutorial.state.operatorstate

import cn.edu.flink.scala.tutorial.source.Event
import cn.edu.flink.scala.tutorial.source.{ClickSource, Event}
import org.apache.flink.api.common.functions.FlatMapFunction
import org.apache.flink.api.common.state.{ListState, ListStateDescriptor}
import org.apache.flink.runtime.state.{FunctionInitializationContext, FunctionSnapshotContext}
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction
import org.apache.flink.streaming.api.functions.sink.SinkFunction
import org.apache.flink.streaming.api.scala._
import org.apache.flink.util.Collector

import scala.collection.mutable.ListBuffer

object ListStateTest {
  def main(args: Array[String]): Unit = {
    val env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI()

    val clickStream = env.addSource(new ClickSource).assignAscendingTimestamps(_.timestamp)

    clickStream.flatMap(new BufferingOut(10)).print()

    env.execute("operatorstate.ListStateTest")
  }

  class BufferingOut(threshold: Int) extends FlatMapFunction[Event, Event] with CheckpointedFunction {

    private var checkpointState: ListState[Event] = _
    private val bufferedElements: ListBuffer[Event] = ListBuffer[Event]()

    override def flatMap(value: Event, out: Collector[Event]): Unit = {
      bufferedElements.append(value)

      if (bufferedElements.size == threshold) {
        for (element <- bufferedElements) {
          out.collect(element)
        }
        bufferedElements.clear()
      }
    }

    override def snapshotState(context: FunctionSnapshotContext): Unit = {
      checkpointState.clear()

      for (element <- bufferedElements) {
        checkpointState.add(element)
      }
    }

    override def initializeState(context: FunctionInitializationContext): Unit = {
      val checkpointStateDescriptor = new ListStateDescriptor[Event]("checkpointState", classOf[Event])

      checkpointState = context.getOperatorStateStore.getListState(checkpointStateDescriptor)

      if (context.isRestored) {
        import scala.collection.JavaConversions._

        for (element <- checkpointState.get()){
          bufferedElements.append(element)
        }
      }
    }


  }
}
